from flask import Flask, request, Response, stream_with_context
import requests
import logging
from logging.handlers import RotatingFileHandler

app = Flask(__name__)

TELEGRAPH_URL = "https://api.openai.com"

# Configure logging
handler = RotatingFileHandler('proxy.log', maxBytes=10000, backupCount=1)
handler.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
app.logger.addHandler(handler)
app.logger.setLevel(logging.INFO)


@app.route("/", defaults={"path": ""}, methods=["GET", "POST", "PUT", "DELETE"])
@app.route("/<path:path>", methods=["GET", "POST", "PUT", "DELETE"])
def proxy(path):
    global TELEGRAPH_URL
    parts = path.split("/")
    new_path = "/" + "/".join(parts[1:])  # parts[0] is empty and parts[1] is 'openai_proxy'
    url = TELEGRAPH_URL + "/" + new_path
    headers = dict(request.headers)
    headers["Host"] = TELEGRAPH_URL.replace("https://", "")
    headers["Access-Control-Allow-Origin"] = headers.get("Access-Control-Allow-Origin") or "*"
    # 删除空字符串键
    if "User" in headers:
        del headers["User"]
        del headers["Accept-Encoding"]


    # Log the incoming request details
    app.logger.info(f"Request method: {request.method}")
    app.logger.info(f"Request URL: {url}")
    app.logger.info(f"Request headers: {headers}")
    app.logger.info(f"Request data: {request.get_data()}")

    response = requests.request(
        method=request.method,
        url=url,
        headers=headers,
        data=request.get_data(),
        cookies=request.cookies,
        allow_redirects=False,
        stream=True,
    )

    def generate():
        for chunk in response.iter_content(chunk_size=1024):
            if chunk:
                yield chunk

    # Filter out headers not to be forwarded
    excluded_headers = ["content-length", "transfer-encoding", "connection","content-encoding"]
    response_headers = [
        (name, value)
        for (name, value) in response.raw.headers.items()
        if name.lower() not in excluded_headers
    ]

        # 删除空字符串键
    # if "Content-Encoding" in headers:
    #     del headers["Content-Encoding"]


    # Log the response details
    app.logger.info(f"Response status: {response.status_code}")
    app.logger.info(f"Response headers: {response_headers}")
    return Response(stream_with_context(generate()), response.status_code, response_headers)


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=9595)
